#include <stdio.h>
#include <cuda_runtime.h>
#include <cublas_v2.h>

int main() {
    // 初始化 cuBLAS 句柄
    cublasHandle_t handle;
    cublasCreate(&handle);

    // 定义矩阵 A 和向量 x、y
    int m = 3, n = 2;
    float A[6] = {1.0, 2.0, 3.0, 4.0, 5.0, 6.0}; // 3x2 矩阵，行主序存储
    float x[2] = {1.0, 2.0}; // 2 维向量
    float y[3] = {0.0, 0.0, 0.0}; // 3 维向量，初始值为 0

    // 标量
    float alpha = 1.0f;
    float beta = 0.0f;

    // 分配 GPU 内存
    float *d_A, *d_x, *d_y;
    cudaMalloc(&d_A, sizeof(A));
    cudaMalloc(&d_x, sizeof(x));
    cudaMalloc(&d_y, sizeof(y));

    // 复制数据到 GPU
    cudaMemcpy(d_A, A, sizeof(A), cudaMemcpyHostToDevice);
    cudaMemcpy(d_x, x, sizeof(x), cudaMemcpyHostToDevice);
    cudaMemset(d_y, 0, sizeof(y));  // 确保 y 初始值为 0

    // 调用 cublasSgemv
    // A 是行主序存储的，所以要使用 CUBLAS_OP_T（转置）
    cublasSgemv(handle, CUBLAS_OP_T, n, m, &alpha, d_A, n, d_x, 1, &beta, d_y, 1);

    // 复制结果回 CPU
    cudaMemcpy(y, d_y, sizeof(y), cudaMemcpyDeviceToHost);

    // 输出结果
    printf("Result vector y:\n");
    for (int i = 0; i < m; i++) {
        printf("%f\n", y[i]);
    }

    // 释放 GPU 内存
    cudaFree(d_A);
    cudaFree(d_x);
    cudaFree(d_y);

    // 销毁 cuBLAS 句柄
    cublasDestroy(handle);

    return 0;
}
